import torch
import torch.nn as nn

class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=True)
        self.conv2 = nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1, bias=True)
        self.conv3 = nn.Conv2d(128, 64, kernel_size=3, stride=1, padding=1, bias=True)
        self.relu = nn.ReLU()
    
    def forward(self, x):
        conv_out = self.conv1(x)
        conv_out = self.conv2(conv_out)
        conv_out = self.conv3(conv_out)
        relu_out = self.relu(conv_out)
        # 添加一个冗余的 Add 节点用于后续演示删除操作
        add_out = relu_out + torch.ones_like(relu_out)
        return add_out

# 导出 ONNX 模型
model = SimpleModel()
dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(
    model,
    dummy_input,
    "original_model.onnx",
    input_names=["input"],
    output_names=["output"],
    opset_version=13
)

print("基础模型已生成：original_model.onnx")